/*
 * Copyright (c) 2024 iSOFT INFRASTRUCTURE SOFTWARE CO., LTD.
 * easyAda is licensed under Mulan PubL v2.
 * You can use this software according to the terms and conditions of the Mulan PubL v2.
 * You may obtain a copy of Mulan PubL v2 at:
 *          http://license.coscl.org.cn/MulanPubL-2.0
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PubL v2 for more details.
 */

#include <tools/rbtree.h>

enum rb_node_color { COLOR_RED, COLOR_BLACK };

enum rb_node_insert_position { INSERT_LEFT, INSERT_RIGHT };

#define NILL ((void *)0)

static inline void initial_node(rbnode_t *node, rbnode_t *parent, unsigned color)
{
    node->parent = parent;
    node->color  = color;
    node->lchild = NILL;
    node->rchild = NILL;
}

static void rotate_left(rbnode_t *tree, rbnode_t *center)
{
    rbnode_t *parent = center->parent;
    rbnode_t *rc     = center->rchild;

    if (parent == tree) {
        parent->parent = rc;
    } else if (parent->lchild == center) {
        parent->lchild = rc;
    } else {
        parent->rchild = rc;
    }

    rc->parent     = parent;
    center->rchild = rc->lchild;
    if (center->rchild != NILL) {
        center->rchild->parent = center;
    }
    rc->lchild     = center;
    center->parent = rc;
}

static void rotate_right(rbnode_t *tree, rbnode_t *center)
{
    rbnode_t *parent = center->parent;
    rbnode_t *lc     = center->lchild;

    if (parent == tree) {
        parent->parent = lc;
    } else if (parent->lchild == center) {
        parent->lchild = lc;
    } else {
        parent->rchild = lc;
    }

    lc->parent     = parent;
    center->lchild = lc->rchild;
    if (center->lchild != NILL) {
        center->lchild->parent = center;
    }
    lc->rchild     = center;
    center->parent = lc;
}

static void insert_fixup(rbnode_t *tree, rbnode_t *node)
{
    rbnode_t *parent;
    rbnode_t *uncle;
    rbnode_t *gparent;
    int       nodeside;
    int       parentside;

    while (node != tree) {
        parent = node->parent;
        if (parent->color == COLOR_BLACK) {
            if (parent == tree) {
                node->color = COLOR_BLACK;
            }
            break;
        }

        nodeside = (parent->lchild == node) ? INSERT_LEFT : INSERT_RIGHT;

        gparent = parent->parent;
        if (gparent->lchild == parent) {
            parentside = INSERT_LEFT;
            uncle      = gparent->rchild;
        } else {
            parentside = INSERT_RIGHT;
            uncle      = gparent->lchild;
        }

        if ((uncle != NILL) && (uncle->color == COLOR_RED)) {
            parent->color  = COLOR_BLACK;
            uncle->color   = COLOR_BLACK;
            gparent->color = COLOR_RED;

            node = gparent;
            continue;
        } else {
            if (parentside != nodeside) {
                if (nodeside == INSERT_LEFT) {
                    rotate_right(tree, parent);
                } else {
                    rotate_left(tree, parent);
                }

                parent = node;
            }

            parent->color  = COLOR_BLACK;
            gparent->color = COLOR_RED;
            if (parentside == INSERT_LEFT) {
                rotate_right(tree, gparent);
            } else {
                rotate_left(tree, gparent);
            }

            break;
        }
    }
}

void rbtree_insert(rbnode_t *tree, rbnode_t *node, int (*compare_less)(rbnode_t *lhs, rbnode_t *rhs))
{
    rbnode_t *parent = tree;
    rbnode_t *pos    = tree->parent;

    while (pos != NILL) {
        parent = pos;

        if (compare_less(node, pos) != 0) {
            pos = pos->lchild;
        } else {
            pos = pos->rchild;
        }
    }

    if (parent == tree) {
        initial_node(node, parent, COLOR_BLACK);
        tree->parent = node;
        return;
    }

    initial_node(node, parent, COLOR_RED);
    if (compare_less(node, parent) != 0) {
        parent->lchild = node;
    } else {
        parent->rchild = node;
    }

    insert_fixup(tree, node);
}

static void delete_fixup(rbnode_t *tree, rbnode_t *node, rbnode_t *parent);
static void delete_fixup_nill(rbnode_t *tree, rbnode_t *node, rbnode_t *parent)
{
    rbnode_t *brother;
    int       nillside = INSERT_LEFT;

    if (parent->lchild == node) {
        brother = parent->rchild;
    } else {
        nillside = INSERT_RIGHT;
        brother  = parent->lchild;
    }

    if (brother->color == COLOR_RED) {
        if (nillside == INSERT_LEFT) {
            parent->lchild         = brother->lchild;
            brother->lchild        = NILL;
            parent->lchild->parent = parent;

            parent  = brother;
            brother = brother->rchild;
        } else {
            parent->rchild         = brother->rchild;
            brother->rchild        = NILL;
            parent->rchild->parent = parent;

            parent  = brother;
            brother = brother->lchild;
        }
    }

    if ((brother->lchild != NILL) || (brother->rchild != NILL)) {
        if ((nillside == INSERT_LEFT) && (brother->rchild == NILL)) {
            rotate_right(tree, brother);
            brother = brother->parent;
        } else if ((nillside == INSERT_RIGHT) && (brother->lchild == NILL)) {
            rotate_left(tree, brother);
            brother = brother->parent;
        } else {
            if (nillside == INSERT_LEFT) {
                brother->rchild->color = COLOR_BLACK;
            } else {
                brother->lchild->color = COLOR_BLACK;
            }
            brother->color = COLOR_RED;
        }

        if (parent->color == COLOR_RED) {
            if (nillside == INSERT_LEFT) {
                rotate_left(tree, parent);
            } else {
                rotate_right(tree, parent);
            }
            parent->color = COLOR_BLACK;
            return;
        } else {
            if (nillside == INSERT_LEFT) {
                rotate_left(tree, parent);
            } else {
                rotate_right(tree, parent);
            }

            tree->parent->color = COLOR_BLACK;
            return;
        }
    } else {
        if (parent->color == COLOR_RED) {
            parent->color  = COLOR_BLACK;
            brother->color = COLOR_RED;
            return;
        } else {
            brother->color = COLOR_RED;
            delete_fixup(tree, parent, parent->parent);
            return;
        }
    }
}

static void delete_fixup(rbnode_t *tree, rbnode_t *node, rbnode_t *parent)
{
    rbnode_t *brother;

    if ((node == NILL) && (parent != tree)) {
        delete_fixup_nill(tree, node, parent);
        return;
    }

    while ((node != tree->parent) && (node != NILL) && (node->color == COLOR_BLACK)) {
        parent  = node->parent;
        brother = (parent->lchild == node) ? parent->rchild : parent->lchild;

        if (brother->color == COLOR_RED) {
            parent->color  = COLOR_RED;
            brother->color = COLOR_BLACK;

            if (parent->lchild == node) {
                rotate_left(tree, parent);
            } else {
                rotate_right(tree, parent);
            }

            brother = parent->rchild;
        }

        if ((brother->lchild == NILL || brother->lchild->color == COLOR_BLACK) &&
            (brother->rchild == NILL || brother->rchild->color == COLOR_BLACK)) {
            brother->color = COLOR_RED;
            node           = parent;
        } else {
            if (parent->lchild == node) {
                if (brother->lchild != NILL && brother->lchild->color == COLOR_RED &&
                    (brother->rchild == NILL || brother->rchild->color == COLOR_BLACK)) {
                    rotate_right(tree, brother);
                    brother = parent->rchild;
                }

                if ((brother->rchild != NILL) && (brother->rchild->color == COLOR_RED)) {
                    brother->color         = parent->color;
                    parent->color          = COLOR_BLACK;
                    brother->rchild->color = COLOR_BLACK;
                    rotate_left(tree, parent);
                    node = tree->parent;
                }
            } else {
                if (brother->rchild != NILL && brother->rchild->color == COLOR_RED &&
                    (brother->lchild == NILL || brother->lchild->color == COLOR_BLACK)) {
                    rotate_left(tree, brother);
                    brother = parent->lchild;
                }

                if ((brother->lchild != NILL) && (brother->lchild->color == COLOR_RED)) {
                    brother->color         = parent->color;
                    parent->color          = COLOR_BLACK;
                    brother->lchild->color = COLOR_BLACK;
                    rotate_right(tree, parent);
                    node = tree->parent;
                }
            }
        }
    }

    node->color = COLOR_BLACK;
}

void rbtree_delete(rbnode_t *tree, rbnode_t *node)
{
    if (node == NILL) {
        return;
    }

    if ((node->lchild == NILL) || (node->rchild == NILL)) {
        rbnode_t *child  = (node->lchild == NILL) ? node->rchild : node->lchild;
        rbnode_t *parent = node->parent;

        if (parent == tree) {
            parent->parent = child;
            if (child != NILL) {
                child->parent = tree;
                child->color  = COLOR_BLACK;
            }
            return;
        } else if (parent->lchild == node) {
            parent->lchild = child;
        } else {
            parent->rchild = child;
        }

        if (child != NILL) {
            child->parent = parent;
        }

        if (node->color != COLOR_RED) {
            delete_fixup(tree, child, parent);
        }
    } else {
        rbnode_t *target = rbnode_next(tree, node);
        rbnode_t *tmp;
        unsigned  color;

        tmp            = node->lchild;
        node->lchild   = target->lchild;
        target->lchild = tmp;
        tmp->parent    = target;

        tmp          = node->rchild;
        node->rchild = target->rchild;
        if (node->rchild != NILL) {
            node->rchild->parent = node;
        }
        if (tmp == target) {
            tmp = node;
        }
        target->rchild = tmp;

        tmp                    = node->parent;
        target->rchild->parent = target;

        if (target->parent == node) {
            node->parent = target;
        } else {
            node->parent         = target->parent;
            node->parent->lchild = node;
        }
        target->parent = tmp;
        if (tmp == tree) {
            tmp->parent = target;
        } else if (tmp->lchild == node) {
            tmp->lchild = target;
        } else {
            tmp->rchild = target;
        }

        color         = target->color;
        target->color = node->color;
        node->color   = color;

        rbtree_delete(tree, node);
    }
}

rbnode_t *rbtree_find(const rbnode_t *tree, unsigned long data, int (*compare)(rbnode_t *lhs, unsigned long data))
{
    rbnode_t *cur = tree->parent;

    while (cur != NILL) {
        int result = compare(cur, data);

        if (result == 0) {
            return cur;
        } else if (result < 0) {
            cur = cur->rchild;
        } else {
            cur = cur->lchild;
        }
    }

    return 0;
}

rbnode_t *rbtree_left_most(const rbnode_t *tree)
{
    rbnode_t *cur = tree->parent;

    if (cur == NILL) {
        return 0;
    }

    while (cur->lchild != NILL) {
        cur = cur->lchild;
    }

    return (cur == NILL) ? 0 : cur;
}

rbnode_t *rbnode_next(const rbnode_t *tree, const rbnode_t *node)
{
    if (node->rchild == NILL) {
        while (1) {
            rbnode_t *p = node->parent;

            if (p == tree) {
                return 0;
            } else {
                if (p->lchild == node) {
                    return p;
                } else {
                    node = p;
                }
            }
        }
    }

    node = node->rchild;

    while (node->lchild != NILL) {
        node = node->lchild;
    }

    return (rbnode_t *)node;
}

void rbtree_head_init(rbnode_t *head)
{
    initial_node(head, NILL, COLOR_BLACK);
}

#ifdef RBTREE_TEST_TOOL
static int is_valid_tree(rbnode_t *parent, rbnode_t *root, int *depth, int *nodes)
{
    int ldepth = 1, lnodes = 0;
    int rdepth = 0, rnodes = 0;

    if (root == NILL) {
        *nodes = 0;
        *depth = 0;
        return 1 == 1;
    }

    if (parent->color == COLOR_RED && root->color == COLOR_RED) {
        return 0 == 1;
    }

    if (!is_valid_tree(root, root->lchild, &ldepth, &lnodes)) {
        return 0 == 1;
    }

    if (!is_valid_tree(root, root->rchild, &rdepth, &rnodes)) {
        return 0 == 1;
    }

    if (root->color == COLOR_BLACK) {
        *depth = ldepth + 1;
    } else {
        *depth = ldepth;
    }

    *nodes = rnodes + lnodes + 1;

    return ldepth == rdepth;
}

int tree_check(rbnode_t *tree, int *depth, int *nodes)
{
    if (tree->parent == NILL) {
        *depth = 0;
        *nodes = 0;
        return tree->color == COLOR_BLACK;
    }

    return is_valid_tree(tree, tree->parent, depth, nodes);
}
#endif
